import argparse
import logging
import sys
import traceback

import cv2

import updater
from lite import FaceRecognitionLite


def get_out_fmt(cap, out_path, fps):
    return None if out_path is None else cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (
        int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))), True)


def core(cap, out, show_flag, write_flag):
    lite = FaceRecognitionLite(img=None, do_option=True, debug_mode=True)
    while cap.isOpened():
        try:
            ret, frame = cap.read()
            persons, genders, ages, traffic, img = lite.new_frame(frame).lite()
            logging.info(
                f'read lite: persons: {persons}, genders: {genders}, ages: {ages}, traffic: {traffic}')
            if show_flag:
                cv2.imshow("frame", img)
            if write_flag:
                out.write(img)
            if cv2.waitKey(25) & 0xFF == ord('q'):
                break
        except Exception as e:
            traceback.print_exc(limit=5)
            logging.error(f'error: {e.args[0]}')
            break

    cap.release()
    if out is not None:
        out.release()
    cv2.destroyAllWindows()


def init(arg):
    updater.check_file(str(args.ds_ip), str(args.bed_port), is_test=arg.is_test)
    cap_file = str(arg.cap_file_path)
    if cap_file == "0":
        cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        fps = 6
        show_flag = True
    else:
        cap = cv2.VideoCapture(cap_file)
        fps = cap.get(cv2.CAP_PROP_FPS)
        show_flag = False
    out_path = str(arg.out_file_path).strip()
    out = None
    write_flag = False
    if len(out_path) > 0:
        if len(out_path) > 4 and out_path[-3:] == "mp4":
            write_flag = True
            out = get_out_fmt(cap, out_path, fps)
        else:
            raise RuntimeError("invalid file suffix")
    core(cap, out, show_flag, write_flag)


def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('cap_file_path', type=str, default="0",
                        help='Video file path to parse.')
    parser.add_argument('--out_file_path', type=str, default="",
                        help='Video file path to write.')
    parser.add_argument('--ds_ip', type=str, default="localhost",
                        help='IP address to check file.')
    parser.add_argument('--bed_port', type=str, default="5000",
                        help='IP address to check file.')
    parser.add_argument('--is_test', action="store_true",
                        help='Whether to run test mode or not.')
    return parser.parse_args(argv)


if __name__ == '__main__':
    logging.basicConfig(
        level='INFO',
        format='[%(asctime)s] <%(levelname)s> %(name)s: %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')
    args = parse_arguments(sys.argv[1:])
    init(args)
